#coding:utf8
from __future__ import division
import numpy as np
import matplotlib.pyplot as plt
from scipy import signal
import time
import theano
import theano.tensor as T
from utils import generateData, generate16Data, generate8Data, calc_distance, psk8, psk16
np.random.seed(123)

M,TT,dB,L = 30000, 20000, 20, 12
EqD = int(round((L+10)/2))
SNR = range(-10, 20)

title_size = 18
label_size = 16


# print plt.rcParams.keys()
plt.rcParams['font.sans-serif']=['simhei'] #用来正常显示中文标签

plt.rcParams['axes.unicode_minus']=False #用来正常显示负号
fig = plt.figure()

class Layer(object):
    def __init__(self, inputs, in_size, out_size, activation_function=None):
        self.W = theano.shared(np.random.normal(0,1,(in_size, out_size)))
        self.Wx_plus_b = T.dot(inputs, self.W)
        self.activation_function = activation_function
        if activation_function:
            self.outputs = self.activation_function(self.Wx_plus_b)
        else:
            self.outputs = self.Wx_plus_b

# determine the inputs
z = T.dmatrix('z')
y = T.dmatrix('y')

def score(pdvalue, Tx, n):
    count = 0
    for i in range(len(pdvalue)-20):
        rp = (pdvalue[i][0], pdvalue[i][1])
        if n == 4 and pdvalue[i][0] * Tx[i].real >= 0 and pdvalue[i][1] * Tx[i].imag >=0:
            count += 1
        if n == 8 and calc_distance(rp, (Tx[i].real, Tx[i].imag)) == min(map(lambda p: calc_distance(rp, p), psk8)):
            count += 1
        if n == 16 and calc_distance(rp, (Tx[i].real, Tx[i].imag)) == min(map(lambda p: calc_distance(rp, p), psk16)):
            count += 1
    return count / (len(pdvalue)-20)

# def score2(pdvalue, Tx):
#     score = 0
#     for i in range(len(pdvalue)):

#         if pdvalue[i][0] * Tx[i].real >=0 and pdvalue[i][1] * Tx[i].imag >=0:
#             score += 1
#     return score/len(pdvalue)



def MLP(X,Tx,db):
    Y = X.T
    Y = np.hstack((np.real(Y),np.imag(Y))) #19990 * 26
    Txlist = np.vstack((np.real(Tx), np.imag(Tx))).T  #19990 * 2
    l1 = Layer(y, 2*L+2, 2, None)  #26入2出

    # loss function 误差函数
    cost = T.mean(T.square(l1.outputs - z))

    # compute the gradients
    gW1 = T.grad(cost, l1.W)
    # apply the gradient descent
    learning_rate = 0.05
    train = theano.function(
        inputs = [y,z],
        outputs = cost,
        updates = [(l1.W, l1.W - learning_rate * gW1)])

    # predict
    predict = theano.function(inputs=[y], outputs=l1.outputs)
    for i in range(201):
        err = train(Y, Txlist)
        if i % 50 == 0:
            pass
            # print(i, time.time()-start, err)

    # testY,testTx,testx = generate8Data(30000,20000,db,L)
    # testY = testY.T
    # testY = np.hstack((np.real(testY),np.imag(testY)))
    # predictValue = predict(testY)
    predictValue = predict(Y)
    return predictValue, score(predictValue, Tx, 16)

    # return predictValue, score(predictValue,testTx, 8)

res = []
for db in range(-10, 30):
    X, Tx, x = generate16Data(30000, 20000, db, L)
    predictValue, accuracy = MLP(X,Tx,db)
    print("db: {} MLP ACCURACY: {}".format(db, accuracy))
    res.append(accuracy)

print(res)



# ax3 = fig.add_subplot(2,2,4)
# predictValue = predictValue.T
# ax3.scatter(predictValue[0], predictValue[1])
# ax3.set_title("MLP", fontsize=title_size)
# ax3.set_xlabel(u"实部", fontsize=label_size)
# ax3.set_ylabel(u"虚部", fontsize=label_size)

# ax4 = fig.add_subplot(2,2,1)
# ax4.scatter([-1,-1,1,1],[-1,1,-1,1])
# ax4.set_title("发送符号", fontsize=title_size)
# ax4.set_xlabel(u"实部",fontsize=label_size)
# ax4.set_ylabel(u"虚部", fontsize=label_size)

# # plt.savefig('foo.png')
# # plt.show()
# plt.show()
